package org.flow_limit;

public class TokenBucket {

    Long addTokenPerSeconds;//每秒生成多少个令牌
    Integer n;// 令牌桶阈值，大于该值的请求直接拒绝
    Long currentTokenNum;
    Long nextTokenTimestamp;//下一次令牌应该生成的时间

    public TokenBucket() {
        this(100,10);
    }

    public TokenBucket(int n, long addTokenPerSeconds) {
        this.n = n;
        this.addTokenPerSeconds = addTokenPerSeconds;
        this.currentTokenNum = 0L;
        this.nextTokenTimestamp = System.currentTimeMillis();
    }

    public Boolean get(int need) {
        long sysTime = System.currentTimeMillis();
        refresh(sysTime);
        if (currentTokenNum >= need) {
            currentTokenNum -= need;
            return true;
        }
        return false;
    }

    public void refresh(long currentTimestamp) {
        if (this.nextTokenTimestamp > currentTimestamp) {
            return;
        }
        this.currentTokenNum = Math.min(this.n, this.currentTokenNum + calculateNeedAddTokenNum(currentTimestamp));
        this.nextTokenTimestamp = calculateNextRefreshTime(currentTimestamp);
    }

    private Long calculateNextRefreshTime(long currentTimestamp) {
        long timeAdd = calculateNeedAddTokenNum(currentTimestamp) * addTokenPerSeconds * 1000L;
        return nextTokenTimestamp + timeAdd;
    }

    private Long calculateNeedAddTokenNum(long currentTimestamp) {
        return ((currentTimestamp - nextTokenTimestamp) / addTokenPerSeconds * 1000L) + 1L;
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucket tokenBucket = new TokenBucket();
        for (int i = 0; i < 10; i++) {
            System.out.println(tokenBucket.get(2));
            Thread.sleep(500);
        }

    }
}